#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/12/10 15:05
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : traffic_flow.py
# @Note    : this is note

import cv2.cv2 as cv


class TrafficFlowDetected:
    def __init__(self):
        self.video = "./data/test.mp4"

    def count_detected(self, video=None):
        count = 0
        if video is None:
            video = self.video
        vc = cv.VideoCapture(video)
        writer1 = cv.VideoWriter('myresult.avi',
                                 cv.VideoWriter_fourcc(*'DIVX'),
                                 30,
                                 (1920, 1080),
                                 True
                                 )
        # 窗口排列
        cv.namedWindow("origin", 0)
        cv.resizeWindow("origin", 480, 360)
        cv.moveWindow("origin", 0, 0)

        cv.namedWindow("gray", 0)
        cv.resizeWindow("gray", 480, 360)
        cv.moveWindow("gray", 480 * 1, 0)

        cv.namedWindow("fgmask", 0)
        cv.resizeWindow("fgmask", 480, 360)
        cv.moveWindow("fgmask", 480 * 2, 0)

        cv.namedWindow("median", 0)
        cv.resizeWindow("median", 480, 360)
        cv.moveWindow("median", 480 * 0, 30 + 360 * 1)

        cv.namedWindow("morphology", 0)
        cv.resizeWindow("morphology", 480, 360)
        cv.moveWindow("morphology", 480 * 1, 30 + 360 * 1)

        cv.namedWindow("Contours", 0)
        cv.resizeWindow("Contours", 480, 360)
        cv.moveWindow("Contours", 480 * 2, 30 + 360 * 1)

        # 背景差分法（混合高斯建模）创建一个对象
        bs = cv.createBackgroundSubtractorMOG2(detectShadows=False)

        while vc.isOpened():
            ret, frame = vc.read()
            # ROI = frame[200:900,300,900]
            # cv.line(frame, (0,600), (19200,600), (0,255,0), 2, 4)
            # 原图
            origin = frame
            cv.imshow("origin", origin)
            # 转为灰度图
            gray = cv.cvtColor(origin, cv.COLOR_BGR2GRAY)
            cv.imshow("gray", gray)
            height = gray.shape[0]
            width = gray.shape[1]
            print(width)
            print(height)
            # cv.imshow(gray)
            # ret,Binary = cv.threshold(gary,50,255,cv.THRESH_BINARY)
            # 根据背景差分法提取前景腌膜
            fgmask = bs.apply(gray)
            cv.imshow("fgmask", fgmask)
            # 中值滤波
            median = cv.medianBlur(fgmask, 5)
            cv.imshow("median", median)
            # 形态学操作
            element = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
            element2 = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
            # 开运算
            image = cv.morphologyEx(median, cv.MORPH_OPEN, element)
            # 膨胀运算
            morphology = cv.dilate(image, element2)
            cv.imshow('morphology', morphology)
            # 二值图像提取轮廓，轮廓跟踪算法（Suzuki，1985）
            contours, hierarchy = cv.findContours(morphology, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
            for cnt in contours:
                # 依次取出每一条轮廓，计算点集或灰度图像的非零像素的右上边界矩形
                x, y, w, h = cv.boundingRect(cnt)
                if y + h == 600:
                    count += 1
                # 过滤掉小物体
                if cv.contourArea(cnt) < width * height * 0.0005 or w < width * 0.01 or h < height * 0.01:
                    continue
                cv.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 3)
                writer1.write(frame)
                cv.putText(frame, "Count:" + str(count), (500, 500), cv.FONT_HERSHEY_COMPLEX, 2, (255, 0, 0))
            cv.imshow("Contours", frame)
            k = cv.waitKey(30) & 0xff
            if k == 27:
                break
        vc.release()
        cv.destroyAllWindows()


TrafficFlowDetected().count_detected("./data/test.mp4")

'''

# 转灰度图
gray = cv.cvtColor(origin, cv.COLOR_BGR2GRAY)
# 计算高、宽
height, width = gray.shape[0], gray.shape[1]
# 提取前景腌膜
fgmask = bs.apply(gray)
# 中值滤波
median = cv.medianBlur(fgmask, 5)
# 形态学处理
element = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
element2 = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
# 开运算
image = cv.morphologyEx(median, cv.MORPH_OPEN, element)
# 膨胀运算
morphology = cv.dilate(image, element2)
# 二值图像提取轮廓，轮廓跟踪算法（Suzuki，1985）
contours, hierarchy = cv.findContours(morphology, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
for cnt in contours:
    # 依次取出每一条轮廓，计算点集或灰度图像的非零像素的右上边界矩形
    x, y, w, h = cv.boundingRect(cnt)
    if y + h == 600:
        count += 1
    if cv.contourArea(cnt) < width * height * 0.0005 or w < width * 0.01 or h < height * 0.01:
        continue
    cv.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 3)

'''
